#! /usr/bin/env python3
from pyspark.ml.feature import SQLTransformer
from pyspark.ml.param import Params
from pyspark.ml.param.shared import HasInputCol, HasOutputCol, Param
from pyspark.sql import SparkSession, DataFrame

from gai.v2.spark.base import SparkTransformer


class LookupMapper(SparkTransformer, HasInputCol, HasOutputCol):
    """Given a lookup table specifying tha mapping from a source column
    to a destination column, LookupMapper can transform ``inputCol`` to ``outputCol``
    using the same mapping.


    Args:
        inputCol: the name of the input column
        outputCol: the name of the output column
        lookupTable: the name of the lookup table
        keyCol: the name of the key column in the lookup table
        valueCol: the name of the value column in the lookup table


    >>> spark = SparkSession.builder.master("local").enableHiveSupport().getOrCreate()
    >>> lookup_data = spark.createDataFrame(
    ...     [("com.tencent.mobileqq", "QQ"), ("com.tencent.mm", "Wechat")],
    ...     schema=("pkg_name", "app_name"))
    >>> lookup_data.registerTempTable("pkg_dim_table")
    >>> sample_data = spark.createDataFrame(
    ...     [("Alice", "com.tencent.mobileqq"),
    ...      ("Bob", "com.tencent.mm"),
    ...      ("Clark", "com.sina.weibo")],
    ...     schema=("user", "pkg"))
    >>> lookup_mapper = LookupMapper(inputCol="pkg", outputCol="app",
    ...                              lookupTable="pkg_dim_table", keyCol="pkg_name",
    ...                              valueCol="app_name")
    >>> output_data = lookup_mapper.transform(sample_data)
    >>> output_data.show()
    +-----+--------------------+------+
    | user|                 pkg|   app|
    +-----+--------------------+------+
    |Alice|com.tencent.mobileqq|    QQ|
    |  Bob|      com.tencent.mm|Wechat|
    |Clark|      com.sina.weibo|  null|
    +-----+--------------------+------+
    <BLANKLINE>
    """

    lookupTable = Param(Params._dummy(), "lookupTable", "the lookup table name")
    keyCol = Param(Params._dummy(), "keyCol", "the key column name")
    valueCol = Param(Params._dummy(), "valueCol", "the value column name")

    def __init__(self, inputCol, outputCol, lookupTable, keyCol, valueCol):
        super(LookupMapper, self).__init__()

        self.setInputCol(inputCol) \
            .setOutputCol(outputCol) \
            .setLookupTable(lookupTable) \
            .setKeyCol(keyCol) \
            .setValueCol(valueCol)

    def setLookupTable(self, lookupTable):
        """Sets the lookup table.

        Args:
            lookupTable: the name of the lookup table

        Returns:
            ``self``
        """
        self._paramMap[self.lookupTable] = lookupTable
        return self

    def getLookupTable(self):
        """
        Returns:
            the name of lookup table
        """
        return self.getOrDefault(self.lookupTable)

    def setKeyCol(self, keyCol):
        """Sets the key column.

        Args:
            keyCol: the name of the key column

        Returns:
            ``self``
        """
        self._paramMap[self.keyCol] = keyCol
        return self

    def getKeyCol(self):
        """
        Returns:
            the name of key column
        """
        return self.getOrDefault(self.keyCol)

    def setValueCol(self, valueCol):
        """Sets the value column.

        Args:
            valueCol: the name of the value column

        Returns:
            ``self``
        """
        self._paramMap[self.valueCol] = valueCol
        return self

    def getValueCol(self):
        """
        Returns:
            the name of value column
        """
        return self.getOrDefault(self.valueCol)

    def getInputCols(self):
        """

        Returns:
            the names of input columns
        """
        return [self.getInputCol()]

    def getOutputCols(self):
        """

        Returns:
            the names of output columns
        """
        return [self.getOutputCol()]

    @property
    def _column_appending(self):
        return True

    def _transform(self, dataset: DataFrame):
        stmt = """
            select a.*, b.{destCol} as {outputCol}
            from __THIS__ a left join {lookupTable} b
            on a.{inputCol} = b.{sourceCol}
        """.format(
            inputCol=self.getInputCol(),
            outputCol=self.getOutputCol(),
            lookupTable=self.getLookupTable(),
            sourceCol=self.getKeyCol(),
            destCol=self.getValueCol()
        )

        dataset = SQLTransformer(statement=stmt).transform(dataset)
        return dataset
